library(extrafont)
library(showtext)
showtext_auto(enable = TRUE)
font_add('KaiTi', 'KaiTi.ttf')
marginaleffect_logit <- function (ModelResults, n.sim = 1000, varname, data, val1 = 0, 
          val2 = 1, clusterid) {
       require(arm)
       library(multiwayvcov)
       library(lmtest)
       library(gridExtra)
       library(viridis)
       library(ggridges)
       library(dplyr)
       library(ggplot2)
       library(tidyr)
       library(stringr)
       library(purrr)
       cluster <- data[, clusterid]
       vcov_cluster = lapply(ModelResults, function(x) FUN = cluster.vcov(x, 
                                                                          cluster))
       modSumm = mapply(function(x, y) coeftest(x, y), x = ModelResults, 
                        y = vcov_cluster)
       noModels = length(modSumm)
       set.seed(12345)
       sim = Map(function(x, y) mvrnorm(n = n.sim, coef(x), y), 
                 ModelResults, vcov_cluster)
       fd <- list()
       for (i in 1:length(ModelResults)) {
              X1 <- model.matrix(ModelResults[[i]])
              X2 <- model.matrix(ModelResults[[i]])
              X = model.matrix(ModelResults[[i]])
              value = as.numeric(quantile(X[, varname[i]], c(val1, 
                                                             val2)))
              X1[, varname[i]] = value[1]
              X2[, varname[i]] = value[2]
              fd[[varname[i]]] = apply(apply(X2, 1, function(x) plogis(sim[[i]] %*% 
                                                                              x)) - apply(X1, 1, function(x) plogis(sim[[i]] %*% 
                                                                                                                           x)), 1, mean)
       }
       df_plot <- map(fd, data.frame) %>% map2_df(., names(.), ~mutate(.x, 
                                                                       id = .y))
       names(df_plot)[1] <- "value"
       
       
       df_plot$id <- factor(df_plot$id, levels = varname)
       
       df_plot2 <- df_plot%>% group_by(id) %>%
                     summarise(low = mean(value) - 1.96*sd(value),
                               median = mean(value),
                               high = mean(value) + 1.96*sd(value))%>%
                     mutate(id = as.numeric(id))
       
       p = ggplot(df_plot2) + 
              geom_hline(yintercept = 0, lty = 2) +
              geom_linerange(aes(x = id, ymin = low,
                                 ymax = high),
                             lwd = 1.5, position = position_dodge(width = 1/2)) +
              geom_pointrange(aes(x = id, y =median, ymin =low,
                                  ymax =high),
                              lwd = 1.5, position = position_dodge(width = 1/2),
                              fill = "WHITE") +
              geom_text(aes(x = id + 0.2,
                            y = median,
                            label = format(round(median, 2)))) +
              theme_bw() + xlab("") + ylab("预测概率变化（边际效应）") + 
              theme(legend.position="bottom",
                    legend.title=element_blank(),
                    legend.text = element_text(colour="blue", size=16),
                    axis.text= element_text(size=16),
                    text = element_text(size=16, family = 'KaiTi',face = "bold"),
                    plot.title = element_text(hjust = .5, size = 16))
   
       return(p)
}


marginaleffect_logit2 <- function (ModelResults, n.sim = 1000, varname, data, val1 = 0, 
                                  val2 = 1, clusterid, dv) {
   require(arm)
   library(multiwayvcov)
   library(lmtest)
   library(gridExtra)
   library(viridis)
   library(ggridges)
   library(dplyr)
   library(ggplot2)
   library(tidyr)
   library(stringr)
   library(purrr)
   cluster <- data[, clusterid]
   vcov_cluster = lapply(ModelResults, function(x) FUN = cluster.vcov(x, 
                                                                      cluster))
   modSumm = mapply(function(x, y) coeftest(x, y), x = ModelResults, 
                    y = vcov_cluster)
   noModels = length(modSumm)
   set.seed(12345)
   sim = Map(function(x, y) mvrnorm(n = n.sim, coef(x), y), 
             ModelResults, vcov_cluster)
   fd <- list()
   for (i in 1:length(ModelResults)) {
      X1 <- model.matrix(ModelResults[[i]])
      X2 <- model.matrix(ModelResults[[i]])
      X = model.matrix(ModelResults[[i]])
      value = as.numeric(quantile(X[, varname[i]], c(val1, 
                                                     val2)))
      X1[, varname[i]] = value[1]
      X2[, varname[i]] = value[2]
      fd[[varname[i]]] = apply(apply(X2, 1, function(x) plogis(sim[[i]] %*% 
                                                                  x)) - apply(X1, 1, function(x) plogis(sim[[i]] %*% 
                                                                                                           x)), 1, mean)
   }
   df_plot <- map(fd, data.frame) %>% map2_df(., names(.), ~mutate(.x, 
                                                                   id = .y))
   names(df_plot)[1] <- "value"
   
   
   df_plot$id <- factor(df_plot$id, levels = varname)
   
   df_plot2 <- df_plot%>% group_by(id) %>%
      summarise(low = mean(value) - 1.96*sd(value),
                median = mean(value),
                high = mean(value) + 1.96*sd(value))%>%
      mutate(id = as.numeric(id))
   
   df_plot2$dv <- dv
   
   return(df_plot2)
}



cluster_se <- function (ModelResults, data,  clusterid) {
       require(arm)
       library(multiwayvcov)
       library(lmtest)
       library(gridExtra)
       library(viridis)
       library(ggridges)
       library(dplyr)
       library(ggplot2)
       library(tidyr)
       library(stringr)
       library(purrr)
       cluster <- data[, clusterid]
       vcov_cluster = lapply(ModelResults, function(x) FUN = cluster.vcov(x, cluster))
  
       modSD =  lapply(vcov_cluster, function(x) FUN = sqrt(diag(x)))
       
       return(modSD)
              
}


cluster_pvalue <- function (ModelResults, data,  clusterid) {
  require(arm)
  library(multiwayvcov)
  library(lmtest)
  library(gridExtra)
  library(viridis)
  library(ggridges)
  library(dplyr)
  library(ggplot2)
  library(tidyr)
  library(stringr)
  library(purrr)
  cluster <- data[, clusterid]
  vcov_cluster = lapply(ModelResults, function(x) FUN = cluster.vcov(x, cluster))
  modSD =  lapply(vcov_cluster, function(x) FUN = sqrt(diag(x)))
  
  # p.values <- coeftest(lm1, vcov. = vcov.lm1)[,4]
  # 
   modcoef = Map(function(x, y) coeftest(x, y), x = ModelResults, y = vcov_cluster)
  pvalue = lapply(modcoef, function(x) FUN = x[, "Pr(>|z|)"])
  return(pvalue)
  
}

### for lm coefplot
scaled_coef_lm <- function(ModelResults, data, clusterid, subvar = FALSE, vars = NULL, dv) {
       library(multiwayvcov)
       library(lmtest)
       library(dplyr)
       require(ggplot2)
       library(gridExtra)
       library(tidyr)
       library(stringr)
       library(purrr)
       
       cluster <- data[, clusterid]
       
       vcov_cluster = lapply(ModelResults, function(x) FUN = cluster.vcov(x, 
                                                                          cluster))
       modcoef = Map(function(x, y) coeftest(x, y), x = ModelResults, y = vcov_cluster)
       modcoef <- lapply(modcoef, function(x) FUN =round(x, digits = 4))  
       modelcoef = lapply(modcoef, function(x) FUN = x[, "Estimate"])
       modelse = lapply(modcoef, function(x) FUN = x[, "Std. Error"])
       varnames = lapply(ModelResults, function(x) FUN = names(x$coefficients))
       dfplot <- list()
       for (i in 1:length(modelcoef)){
              ad = 0.5 / modelse[[i]]
              modelcoef[[i]] = modelcoef[[i]]* ad
              modelse[[i]] = modelse[[i]]*ad 
              
              ylo95 =modelcoef[[i]] - qt(.975, nrow(data))*(modelse[[i]])
              yhi95 =modelcoef[[i]] + qt(.975, nrow(data))*(modelse[[i]])
              ylo90 =modelcoef[[i]] - qt(.95, nrow(data))*(modelse[[i]])
              yhi90 =modelcoef[[i]] + qt(.95, nrow(data))*(modelse[[i]])
              dfplot[[i]] = data.frame(varnames[[i]], modelcoef[[i]], modelse[[i]], ylo95, yhi95,
                                       ylo90, yhi90, dv,
                                       Model = paste('Model',seq_along(modelcoef)[i]))
       }
       
       coefficientdata = do.call(rbind,dfplot)
       colnames(coefficientdata) <- c("Variables", "Coefficients", "S.E",
                                      "Low95CI", "High95CI", 
                                      "Low90CI", "High90CI", "DV",
                                      "Model.Name")
       df <- coefficientdata %>%
              dplyr::mutate(Variables = as.character(Variables)) %>%
              dplyr::arrange(Variables)
       
       if(subvar == TRUE){
              df <- df[grep(paste(vars, collapse = "|"), df$Variables),]
              #re-order the variables: if factor, the order of factor levels is used; if character, an alphabetical order ist used
              df$Variables <- factor(df$Variables, levels = vars)
       } 
   
       
       return(df)   
} 


## functions
MLM_coef_df <- function(modResults, data, vars,dv) {
       library(extrafont)
       library(showtext)
       # font_add('SimSun', 'SimSun.ttf')
       showtext_auto(enable = TRUE)
       modelcoef=lapply(modResults, 
                        function(x) FUN=fixef(x))
       modelse = lapply(modResults, 
                        function(x) FUN=sqrt(diag(vcov(x))))
       
       varnames =lapply(modResults, function(x) FUN = names(fixef(x)))
       
       dfplot <- list()
       
       for (i in 1:length(modelcoef)){
              ad = 0.5 / modelse[[i]]
              modelcoef[[i]] = modelcoef[[i]]* ad
              modelse[[i]] = modelse[[i]]*ad 
              
              ylo95 =modelcoef[[i]] - qt(.975, nrow(data))*(modelse[[i]])
              yhi95 =modelcoef[[i]] + qt(.975, nrow(data))*(modelse[[i]])
              ylo90 =modelcoef[[i]] - qt(.95, nrow(data))*(modelse[[i]])
              yhi90 =modelcoef[[i]] + qt(.95, nrow(data))*(modelse[[i]])
              dfplot[[i]] = data.frame(varnames[[i]], modelcoef[[i]], modelse[[i]], ylo95, yhi95,
                                       ylo90, yhi90,dv,
                                       Model = paste('Model',seq_along(modelcoef)[i]))
       }
       coefficientdata = do.call(rbind,dfplot)
       colnames(coefficientdata) <- c("Variables", "Coefficients", "S.E",
                                      "Low95CI", "High95CI", 
                                      "Low90CI", "High90CI", "DV",
                                      "Model.Name")
       
       df <- coefficientdata %>%
              dplyr::mutate(Variables = as.character(Variables)) %>%
              dplyr::arrange(desc(Variables))
       
       df <-  df[grep(paste(vars, collapse = "|"), df$Variables),]
       
       #re-order the variables: if factor, the order of factor levels is used; if character, an alphabetical order ist used
       df$Variables <- factor(df$Variables, levels = vars)
       
       return(df)
}  


plot_interEffect_binary <- function(ModelResults, bvarname1, cvarname2,
                                    cval1, cval2, clusterid, data, dv, iv){
   
   require(arm)
   library(multiwayvcov)
   library(lmtest)
   library(gridExtra)
   library(viridis)
   library(ggridges)
   library(dplyr)
   library(ggplot2)
   library(tidyr)
   library(stringr)
   library(purrr)
   cluster <- data[, clusterid]
   vcov_cluster = cluster.vcov(ModelResults, cluster)
   
   draw <- mvrnorm(n = 1000, coef(ModelResults), vcov_cluster)
   
   ##set simulation
   df <- array(NA, c(2, 4))
   
   X1 <- model.matrix(ModelResults)
   X2 <- model.matrix(ModelResults)
   X1[, bvarname1] = 0
   X1[, cvarname2] = cval1
   X1[, paste(cvarname2, bvarname1, sep = ":")] = 0*cval1
   
   X2[,bvarname1] = 0
   X2[,cvarname2] = cval2
   X2[, paste(cvarname2, bvarname1, sep = ":")] = 0*cval2
   
   fd_no <- apply(apply(X2, 1, function (x) draw %*% x) - 
                     apply(X1, 1, function (x) draw %*% x), 1, mean) #1 indicates row, 2= columns
   
   X3 <- model.matrix(ModelResults)
   X4 <- model.matrix(ModelResults)
   X3[, bvarname1] = 1
   X3[, cvarname2] = cval1
   X3[, paste(cvarname2, bvarname1, sep = ":")] = 1*cval1
   
   X4[,bvarname1] = 1
   X4[,cvarname2] = cval2
   X4[, paste(cvarname2, bvarname1, sep = ":")] = 1*cval2
   
   fd_yes <- apply(apply(X4, 1, function (x) draw %*% x) - 
                      apply(X3, 1, function (x) draw %*% x), 1, mean) #1 indicates row, 2= columns
   
   df[1, 1] <- 0
   df[1, 2] <- mean(fd_no)
   df[1, 3:4] <- quantile(fd_no, probs = c(.05,.95))
   df[2, 1] <- 1
   df[2, 2] <- mean(fd_yes)
   df[2, 3:4] <- quantile(fd_yes, probs = c(.05,.95))
   
   df_plot <- as.data.frame(df)
   
   
   
   colnames(df_plot) <- c("type", "mean", "lo", "hi")   
   df_plot$type <- as.factor(df_plot$type)
   
   df_plot$DV <- dv
   df_plot$IV <- cvarname2
   
   return(df_plot)
}



### for lm coefplot
coef_lm <- function(ModelResults, data, clusterid, subvar = FALSE, vars = NULL, dv) {
   library(multiwayvcov)
   library(lmtest)
   library(dplyr)
   require(ggplot2)
   library(gridExtra)
   library(tidyr)
   library(stringr)
   library(purrr)
   
   cluster <- data[, clusterid]
   
   vcov_cluster = lapply(ModelResults, function(x) FUN = cluster.vcov(x, 
                                                                      cluster))
   modcoef = Map(function(x, y) coeftest(x, y), x = ModelResults, y = vcov_cluster)
   modcoef <- lapply(modcoef, function(x) FUN =round(x, digits = 4))  
   modelcoef = lapply(modcoef, function(x) FUN = x[, "Estimate"])
   modelse = lapply(modcoef, function(x) FUN = x[, "Std. Error"])
   varnames = lapply(ModelResults, function(x) FUN = names(x$coefficients))
   dfplot <- list()
   for (i in 1:length(modelcoef)){
      # ad = 0.5 / modelse[[i]]
      # modelcoef[[i]] = modelcoef[[i]]* ad
      # modelse[[i]] = modelse[[i]]*ad 
      
      ylo95 =modelcoef[[i]] - qt(.975, nrow(data))*(modelse[[i]])
      yhi95 =modelcoef[[i]] + qt(.975, nrow(data))*(modelse[[i]])
      ylo90 =modelcoef[[i]] - qt(.95, nrow(data))*(modelse[[i]])
      yhi90 =modelcoef[[i]] + qt(.95, nrow(data))*(modelse[[i]])
      dfplot[[i]] = data.frame(varnames[[i]], modelcoef[[i]], modelse[[i]], ylo95, yhi95,
                               ylo90, yhi90, dv,
                               Model = paste('Model',seq_along(modelcoef)[i]))
   }
   
   
   coefficientdata = do.call(rbind,dfplot)
   colnames(coefficientdata) <- c("Variables", "Coefficients", "S.E",
                                  "Low95CI", "High95CI", 
                                  "Low90CI", "High90CI", "DV",
                                  "Model.Name")
   df <- coefficientdata %>%
      dplyr::mutate(Variables = as.character(Variables)) %>%
      dplyr::arrange(Variables)
   
   if(subvar == TRUE){
      df <- df[grep(paste(vars, collapse = "|"), df$Variables),]
      #re-order the variables: if factor, the order of factor levels is used; if character, an alphabetical order ist used
      df$Variables <- factor(df$Variables, levels = vars)
   } 
   
   marg <- function(x){x*100}
   
   df <- df %>% dplyr::mutate_at(vars("Coefficients", "S.E",
                                      "Low95CI", "High95CI", 
                                      "Low90CI", "High90CI"), list(mar = ~marg(.) ))
   df$id <- as.numeric(as.factor(df$Variables))
   p =  ggplot(df) + 
      geom_hline(yintercept = 0, lty = 2) +
      geom_linerange(aes(x = Variables, ymin = Low95CI_mar,
                         ymax = High95CI_mar),
                     lwd = 1.5, position = position_dodge(width = 1/2)) +
      geom_pointrange(aes(x = Variables, y =Coefficients_mar, ymin = Low95CI_mar,
                          ymax = High95CI_mar),
                      lwd = 1.5, position = position_dodge(width = 1/2),
                      fill = "WHITE")+
      geom_text(aes(x = id +0.2 ,
                    y =Coefficients_mar,
                    label = format(round(Coefficients_mar, 4))),position = position_dodge(width = 1/2)) +
      theme_bw()
   
   return(p)   
} 



